# mcp_research_assistant.py

import logging
from typing import List, Dict, Any
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Settings,
    PromptTemplate,
    FunctionTool,
)
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.tools.code import CodeInterpreterTool
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import ToolMetadata

# --- 配置 ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Ollama 配置
OLLAMA_BASE_URL = "http://localhost:11434"
LLM_MODEL = "qwen:8b"
EMBEDDING_MODEL = "nomic-embed-text:1.5"

# 本地知识库路径
RESEARCH_NOTES_DIR = "./research_notes"
STORAGE_DIR = "./storage_research_assistant"

# --- 客制化 Prompt 模板 (核心) ---
# 这些 Prompt 将精细控制智能体的行为

# 1. 主系统 Prompt (定义角色、能力、决策逻辑)
CUSTOM_SYSTEM_PROMPT = PromptTemplate("""
你是一位顶尖的 AI 研究助手，代号 "Researcher-X"。
你拥有三种核心能力：
1. **内部知识检索 (internal_knowledge)**: 访问用户本地的研究笔记和文档。
2. **网络探索 (web_search)**: 通过 DuckDuckGo 搜索获取最新的外部信息。
3. **数据计算 (code_execution)**: 执行 Python 代码进行数学计算、数据分析或生成图表。

### 决策协议 (MCP)
- **优先级**: 内部知识 > 网络探索 > 数据计算。
- **步骤**:
  1.  首先检查问题是否能用内部知识回答。如果能，直接回答并引用来源。
  2.  如果问题涉及外部事实、最新进展或公共数据，使用 `web_search`。
  3.  如果问题需要数值计算、公式推导、数据处理或绘图，使用 `code_execution`。
  4.  可以组合使用工具（例如，先搜索数据，再用代码分析）。
- **禁止行为**:
  - 绝不编造信息。
  - 不要同时调用多个工具除非必要。
  - 执行代码前，向用户简要说明将要执行的操作。

### 输出格式
- **回答**: 清晰、准确、简洁。
- **工具调用后**: 先解释结果，再给出最终结论。
- **引用**:
  - 来自内部知识: [来源: {file_name}, 页码 {page_label}]
  - 来自网络搜索: [来源: DuckDuckGo 搜索, {url}]
  - 来自代码执行: [来源: Python 代码执行]
- **无法回答**: "根据我的知识和能力，无法回答该问题。"
""")

# 2. (可选) 工具特定 Prompt - 用于指导代码解释器
CODE_TOOL_PROMPT = PromptTemplate("""
你是一个精确的 Python 代码解释器。
- 仅执行用户要求的计算或数据处理。
- 使用 matplotlib 生成图表时，确保代码完整且能显示。
- 执行后，返回关键结果和观察。
- 不要添加额外的解释性文本，Agent 会处理。
""")

# --- 设置 LLM 和 Embedding ---
Settings.llm = Ollama(
    model=LLM_MODEL,
    base_url=OLLAMA_BASE_URL,
    request_timeout=180.0,  # 工具调用可能需要更长时间
    system_prompt=CUSTOM_SYSTEM_PROMPT.template,  # 应用客制化系统提示
)

Settings.embed_model = OllamaEmbedding(
    model_name=EMBEDDING_MODEL,
    base_url=OLLAMA_BASE_URL,
)


# --- 构建工具集 (MCP Tools) ---
def create_tools() -> List[FunctionTool]:
    tools = []

    # 工具 1: 本地研究笔记知识库
    try:
        if not Path(RESEARCH_NOTES_DIR).exists():
            Path(RESEARCH_NOTES_DIR).mkdir()
            logger.warning(f"目录 '{RESEARCH_NOTES_DIR}' 不存在，已创建。请放入研究笔记。")

        # 加载或创建本地知识库索引
        if Path(STORAGE_DIR).exists():
            storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
            index = load_index_from_storage(storage_context)
        else:
            documents = SimpleDirectoryReader(RESEARCH_NOTES_DIR).load_data()
            index = VectorStoreIndex.from_documents(documents, show_progress=True)
            index.storage_context.persist(persist_dir=STORAGE_DIR)

        query_engine = index.as_query_engine(similarity_top_k=3)
        knowledge_tool = QueryEngineTool(
            query_engine=query_engine,
            metadata=ToolMetadata(
                name="internal_knowledge",
                description="查询用户本地的研究笔记、实验记录和参考文档。包含机器学习、数学、项目计划等内容。",
            ),
        )
        tools.append(knowledge_tool)
        logger.info("✅ 本地知识库工具已加载。")
    except Exception as e:
        logger.error(f"加载本地知识库失败: {e}")

    # 工具 2: 网络搜索
    try:
        search_tool_spec = DuckDuckGoSearchToolSpec()
        # 通常 to_tool_list() 返回一个包含多个工具的列表，我们取搜索工具
        web_search_tool = search_tool_spec.to_tool_list()[0]
        # 可以选择性地修改其元数据
        web_search_tool.metadata = ToolMetadata(
            name="web_search",
            description="通过 DuckDuckGo 搜索引擎查找最新的信息、新闻、定义或公共数据。用于回答涉及外部事实的问题。",
        )
        tools.append(web_search_tool)
        logger.info("✅ 网络搜索工具已加载。")
    except Exception as e:
        logger.error(f"初始化网络搜索工具失败: {e}")

    # 工具 3: 代码执行 (Python)
    try:
        # 创建 CodeInterpreterTool
        code_tool = CodeInterpreterTool(
            # 可以传入客制化 prompt
            system_prompt=CODE_TOOL_PROMPT.template,
            # timeout=60, # 可选
        )
        # 包装成 FunctionTool 并设置元数据
        code_function_tool = FunctionTool.from_defaults(
            fn=code_tool,
            name="code_execution",
            description="执行 Python 代码进行数学计算、数据分析、公式求解或生成图表（使用 matplotlib）。用于处理需要精确计算或可视化的任务。",
        )
        tools.append(code_function_tool)
        logger.info("✅ 代码执行工具已加载。")
    except Exception as e:
        logger.error(f"初始化代码执行工具失败: {e}")

    return tools


# --- 创建 MCP Agent ---
def create_agent() -> ReActAgent:
    tools = create_tools()
    if not tools:
        raise ValueError("未能创建任何工具。")

    agent = ReActAgent.from_tools(
        tools=tools,
        llm=Settings.llm,
        verbose=True,  # 输出详细的推理和工具调用过程
        # 注意：system_prompt 已在 Settings.llm 中全局设置
        # 如果需要覆盖，可以在这里传入
        # system_prompt="..."
    )
    return agent


# --- 主程序 ---
def main():
    print("🚀 MCP AI 研究助手已启动 (qwen:8b + nomic-embed-text:1.5)")
    print("能力: 本地知识检索 | 网络搜索 | 代码执行")
    print("输入 'quit' 退出。")
    print("-" * 60)

    try:
        agent = create_agent()
    except Exception as e:
        logger.error(f"创建智能体失败: {e}")
        return

    while True:
        query = input("\n🔍 您的研究问题: ").strip()
        if query.lower() == "quit":
            print("再见！")
            break
        if not query:
            continue

        try:
            logger.info(f"处理问题: {query}")
            response = agent.chat(query)
            print(f"\n🤖 Researcher-X: {response.response}")

            # (可选) 显示源信息
            # if hasattr(response, 'source_nodes') and response.source_nodes:
            #     print(f"\n📚 信息来源:")
            #     for node in response.source_nodes:
            #         print(f"  - {node.node.metadata.get('file_name', 'Unknown')}")

        except Exception as e:
            print(f"❌ 处理问题时出错: {e}")


if __name__ == "__main__":
    # 必须从 llama_index.core 导入 StorageContext 和 load_index_from_storage
    from llama_index.core import StorageContext, load_index_from_storage
    from pathlib import Path

    main()